# 毕业设计main.py
import os
import math
import argparse
import json
import time
import copy
import numpy as np

import torch
from get_weight_matrix import get_matrix
import matplotlib.pyplot as plt
import torch.utils.data
import torch.optim as optim
from torchvision import transforms
import torch.multiprocessing as mp
import torch.distributed as dist
import torch.optim.lr_scheduler as lr_scheduler

import sys
sys.path.append('extra_utils')
from extra_utils.model import resnet34, resnet101, mnist_Net
from extra_utils.distributed_utils import init_distributed_mode, cleanup, is_main_process
from client import *
import my_data_loader

# 此时不可以正常显示中文标题，则加上以下代码
plt.rcParams['font.sans-serif'] = ['KaiTi', 'FangSong', 'SimHei']
plt.rcParams['font.size'] = 12
plt.rcParams['axes.unicode_minus'] = False


def main(rank, args, conf):
    device = torch.device(args.device)

    if os.path.exists("./weights") is False:  # 判断当前路径下有没有weights文件夹，如果没有，则创建文件夹
        os.makedirs("./weights")  # 用于保存训练权重，方便在使用torch.save时奏效。

    init_distributed_mode(conf, rank)  # 分布式包函数初始化

    train_loader, val_loader, train_sampler, train_length, eval_length =\
        my_data_loader.get_dataloader(args, conf, rank, name=conf['type'])

    # 如果存在预训练权重则载入
    # 通过json文件嵌套字典实现了灵活访问num_classes model_name
    # model = resnet34(num_classes=conf['num_classes'][conf['type']]).to(device)
    torch.manual_seed(1)  # 将我们的初始实例化模型在每个不同的进程组中实现复现。
    if conf['model_name'][conf['type']] == "resnet34":
        model = resnet34(num_classes=conf['num_classes'][conf['type']]).to(device)
    else:
        model = mnist_Net().to(device)

    # 判断 “载入模型权重”
    if args.pre_trained:
        weights_path = conf['initial_weights'][conf['model_name'][conf['type']]]
        if weights_path != "":  # 如果我们传入weights参数
            if os.path.exists(weights_path):  # 判断是否存在权重
                weights_dict = torch.load(weights_path, map_location=device)  # 载入模型权重，返回一个(按层)有序字典
                load_weights_dict = {k: v for k, v in weights_dict.items()  # 遍历每层模型参数
                                     if model.state_dict()[k].numel() == v.numel()}
                # torch.numel()为返回数组中元素的个数，在此判断是否数目一样。如果一样则保存在字典当中，之后载入该层权重，如果不一样则不保存（即不保存该层权重）
                # 这样做的目的是因为预训练模型是处理imageNet的1000分类问题，全连接层和我们的花分类问题不一样，
                model.load_state_dict(load_weights_dict, strict=False)  # 载入权重
            else:
                raise FileNotFoundError("not found weights file: {}".format(weights_path))

    global_model = model
    dist.barrier()

    cost_list = []  # 构建损失函数列表，方便进行数据可视化
    cost_1 = []  # rank1的loss函数列表
    cost_2 = []  # rank2的loss函数列表
    accuracy_list = []  # 构建准确率列表
    accuracy_aggregate = []  # 构建模型聚合后的准确率列表
    accuracy_1 = []  # rank1的准确率列表
    accuracy_2 = []  # rank2的准确率列表
    class_mean_list = []  # 用于存放类平均值的列表

    # 客户端和服务端实例化;rank0既作为server端又作为client端。
    torch.manual_seed(5)
    # 我们使用随机种子成功将三个进程的初始模型设置成为相同，并为后续复现情况提供了技术参考。
    # seed值不同，则每次的采取的随机情况也不同。
    client = Client(args, conf, train_loader, train_sampler, val_loader, rank)
    print("client %d creation complete！" % rank)
    dist.barrier()  # 确保所有进程所对应的客户端client类实例化完成。

    acc, aver_loss = client.model_eval(global_model)  # 先模型评估得到训练前的模型预测的准确率
    dist.barrier()
    print("Rank %d Global Epoch %d, accuracy: %f, loss: %f\n" % (rank, 0, acc, aver_loss))
    accuracy_list.append(acc)
    accuracy_1.append(acc)
    accuracy_2.append(acc)
    accuracy_aggregate.append(acc)
    dist.barrier()

    t_start = time.perf_counter()
    num_class = 0

    # 开始全局迭代
    for global_epoch in range(conf['global_epochs']):

        # 我们在该位置曾使用client的model_eval模型评估，得到了相同的accuracy/loss，这证明我们的初始模型是相同的。
        # 即便两次的模型评估挨着，也会有一些误差(具体原因不清楚)。但我们可以使用torch.manual_seed(seed)来复现，使得loss相同。
        # acc, aver_loss = client.model_eval(global_model)
        # dist.barrier()
        # print("Rank %d Global Epoch %d, accuracy: %f, loss: %f\n" % (rank, global_epoch + 1, acc, aver_loss))
        # dist.barrier()

        local_model = copy.deepcopy(global_model)  # local_model将要用于局部迭代
        diff, class_mean = client.local_train(local_model, args, global_epoch, cost_list, train_length, accuracy_list,
                                              accuracy_1, accuracy_2, rank, cost_1, cost_2)

        class_mean_list.append(class_mean)  # 构造关于每个epoch的类平均值的列表, 方便后续数据可视化作图.(模型聚合之前)
        num_class = class_mean.size()[0]  # class_mean.size()[0]代表着类别数目

        # 实现map_reduce
        for key, value in diff.items():  # value
            dist.all_reduce(value, op=dist.ReduceOp.SUM)
            dist.barrier()

        # 全局模型参数更新,得到新的global_model
        for name, value in global_model.state_dict().items():  # 目前尚未成功进行全局模型的更新！
            update_per_layer = diff[name] * conf["lambda"]  # 为了将all_reduce之后并进行平均化
            value = value.float()  # 解决:result type Float can't be cast to the desired output type Long
            update_per_layer = update_per_layer.float()
            value.add_(update_per_layer)
        if is_main_process():
            print("Global_epoch [%d/%d] model_aggregate has been accomplished!" % (global_epoch+1, conf['global_epochs']))

        acc, aver_loss = client.model_eval(global_model)  # 我们是使用整个测试集数据来进行模型评估的。
        accuracy_aggregate.append(acc)
        dist.barrier()
        print("Rank %d Global Epoch %d, accuracy: %f, loss: %f\n" % (rank, global_epoch+1, acc, aver_loss))
        dist.barrier()

    # 得到关于类平均值的 类-epoch 矩阵.并通过分布式通信包,得到各client的类-epoch 矩阵.
    # 我们采用了深拷贝的方法解决了一直以来解决不了的问题。
    class_epoch_matrix, class_epoch_matrix_1, class_epoch_matrix_2 = get_matrix(class_mean_list, conf, rank)

    # 数据可视化
    if is_main_process():
        plt.figure(1)  # 作出损失函数曲线和各client本地以及聚合模型之后的准确率!

        plt.subplot(1, 2, 1)  # 损失函数图线的绘制
        plt.plot(cost_list, label='Rank0的Loss曲线', color='blue', linewidth=1.0)
        plt.plot(cost_1, label='Rank1的Loss曲线', color='orange', linewidth=1.0)
        plt.plot(cost_2, label='Rank2的Loss曲线', color='green', linewidth=1.0)
        plt.legend(loc="best")
        plt.xlabel('total_epoch*5')
        plt.ylabel('Total_Loss')
        plt.title('[{}]Running_Loss的变化趋势'.format(conf["type"]))

        plt.subplot(1, 2, 2)
        plt.plot(accuracy_list, label='Rank0的测试集准确率', color='blue', linewidth=1.0)
        plt.plot(accuracy_1, label='Rank1的测试集准确率', color='orange', linewidth=1.0)
        plt.plot(accuracy_2, label='Rank2的测试集准确率', color='green', linewidth=1.0)
        plt.scatter([idx*5*conf["local_epochs"] for idx in range(conf["global_epochs"]+1)], accuracy_aggregate, s=12, color='purple')
        plt.plot([idx*5*conf["local_epochs"] for idx in range(conf["global_epochs"]+1)], accuracy_aggregate, color='red',
                 linewidth=1.0, label="聚合后的准确率")
        plt.legend(loc="best")
        plt.xlabel('total_epoch*5')
        plt.ylabel('Accuracy')
        plt.title('[{}]Accuracy的变化趋势'.format(conf["type"]))

        plt.figure(2)  # 作出各个client的按client分子图的类均值图线.

        plt.subplot(1, 3, 1)
        for j in range(num_class):  # 10次循环
            plt.plot((class_epoch_matrix[j] * 1000).tolist(), label='class {} 类均值'.format(j), linewidth=1.0)
            plt.legend(loc="best")
            plt.xlabel('epoch')
            plt.ylabel('mean_value')
            plt.title('Rank 0 mean_value')

        plt.subplot(1, 3, 2)
        for j in range(num_class):  # 10次循环
            plt.plot((class_epoch_matrix_1[j] * 1000).tolist(), label='class {} 类均值'.format(j), linewidth=1.0)
            plt.legend(loc="best")
            plt.xlabel('epoch')
            plt.ylabel('mean_value')
            plt.title('Rank 1 mean_value')

        plt.subplot(1, 3, 3)
        for j in range(num_class):  # 10次循环
            plt.plot((class_epoch_matrix_2[j] * 1000).tolist(), label='class {} 类均值'.format(j), linewidth=1.0)
            plt.legend(loc="best")
            plt.xlabel('epoch')
            plt.ylabel('mean_value')
            plt.title('Rank 2 mean_value')

        plt.figure(3)  # 作出各个client的按class分子图的类均值图线.

        for i in range(num_class):
            plt.subplot(2, 5, i + 1)
            plt.plot((class_epoch_matrix[i] * 1000).tolist(), label='client {} 类均值'.format(0), linewidth=1.0)
            plt.plot((class_epoch_matrix_1[i] * 1000).tolist(), label='client {} 类均值'.format(1), linewidth=1.0)
            plt.plot((class_epoch_matrix_2[i] * 1000).tolist(), label='client {} 类均值'.format(2), linewidth=1.0)
            plt.legend(loc="best")
            plt.xlabel('epoch')
            plt.ylabel('mean_value')
            plt.title("class {}".format(i))

        plt.show()

    t_end = time.perf_counter()
    if is_main_process():
        print("全局迭代轮次 %d 共花费了时间 %f 秒" % (conf['global_epochs'], t_end-t_start))

    cleanup()  # 关闭进程组并去初始化


if __name__ == '__main__':
    parser = argparse.ArgumentParser()  # 创建一个解析器——创建 ArgumentParser() 对象

    # 添加参数——调用 add_argument() 方法添加参数
    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--lr', type=float, default=0.001)  # 初始学习率为0.001
    parser.add_argument('--lrf', type=float, default=0.1)  # 倍率因子为0.1，方便之后使得学习率递减为初始学习率的1/10

    # resnet34 官方权重下载地址
    # https://download.pytorch.org/models/resnet34-333f7ec4.pth
    parser.add_argument('--freeze_layers', type=bool, default=False)  # 是否冻结全连接层之前的所有网络。
    parser.add_argument('--device', default='cpu', help='device id (i.e. 0 or 0,1 or cpu)')  # 我们默认的device是cpu
    parser.add_argument('--nw', default=1, type=int, help='number of workers')
    # 是否启用SyncBatchNorm（同步BN层处理）
    parser.add_argument('--syncBN', type=bool, default=False)  # 同步BN层处理，对训练速度是有影响的，但对准确率有好处。
    parser.add_argument('-c', '--conf', default="./conf.json", type=str, dest='conf')
    parser.add_argument('-p', '--pre_trained', default=True, type=bool,
                        help="This parameter determines whether we want to transfer learning")
    parser.add_argument('--only_0', default=False, type=bool, help="console whether we plot only ran0")

    # 解析参数——使用 parse_args() 解析添加的参数
    opt = parser.parse_args()
    with open(opt.conf, 'r') as f:
        conf = json.load(f)

    size = conf['world_size']
    print("World_size : %d" % size)
    mp.spawn(main,  # 主运行函数
             args=(opt, conf, ),  # 主函数的参数 Expected type 'tuple'。主函数的形参一定要有rank
             nprocs=size,  # 当前节点的进程数
             join=True)  # 加入同一进程池

# 主进程既要有server的功能，又要有client的功能
# 由于pytorch多进程分布式通信包只能发送接收tensor张量。我们本想通过：服务端发送全局初始模型、全局模型到各个client客户端执行本地训练，更新模型参数
# ，求得当前轮次的模型权重差，再由客户端发送模型权重差到服务端，实现联邦平均算法，得到新一轮次的全局模型。
# 于是我们假想方法：按每层依次发送/接受对应层的模型权重差：(dist.all_reduce)
# 依次发送diff字典的每个键值数据，由server收到存放到本地的diff字典中，分别将diff各键对应的值相加(聚合一个新的模型)。再将总的权重差发送回各客户端即
# 相当于更新全局模型。
